from posixpath import join
import requests
import os
import base64
import json
import time
import re


def cv2_to_base64(image):
    return base64.b64encode(image).decode('utf8')

def predict(file_path):
    url = "http://127.0.0.1:9998/ocr/prediction"
    with open(file_path, 'rb') as file:
        image_data1 = file.read()
    image = cv2_to_base64(image_data1)
    data = {"key": ["image"], "value": [image]}
    start_time = time.time()
    r = requests.post(url=url, data=json.dumps(data))
    print(f'耗时：{time.time() - start_time}')
    print(r.json())
    values = r.json().get('value')
    # print(values[0], type(eval(values[0])))
    return_dict = {}
    for item in eval(values[0]):
        if '姓名' in item:
            return_dict['name'] = item.replace('姓名', '')
        elif len(item) == 18 and (item.isnumeric() or item[:-1].isnumeric()):
            return_dict['idNo'] = item
        elif '公民身份号码' in item:
            return_dict['idNo'] = item.replace('公民身份号码', '').replace(' ', '')
        elif '性别' in item:
            return_dict['gender'] = item.replace('性别', '')
        elif '民族' in item:
            return_dict['nation'] = item.replace('民族', '')
        elif '出生' in item:
            return_dict['birthday'] = '-'.join(re.findall( r'\d+', item, re.M|re.I))
        elif '住址' in item:
            return_dict['address'] = item.replace('住址', '')
        elif item == '公民身份号码':
            continue
        else:
            address = return_dict.get('address', '')
            address += item
            return_dict['address'] = address
    return return_dict
    
def main():
    base_dir = os.path.join(os.getcwd(), 'test1') 
    files = os.listdir(base_dir)
    for item in files:
        file_path = os.path.join(base_dir, item)
        res = predict(file_path)
        print(res)

if __name__ == '__main__':
    main()